package net.minecraft.server.world;

import com.mojang.datafixers.util.Pair;
import com.mojang.logging.LogUtils;
import com.mojang.serialization.Codec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import it.unimi.dsi.fastutil.longs.Long2ObjectFunction;
import it.unimi.dsi.fastutil.longs.Long2ObjectMaps;
import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
import it.unimi.dsi.fastutil.longs.LongSet;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap.Entry;
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import it.unimi.dsi.fastutil.objects.ObjectIterator;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.function.BiConsumer;
import java.util.function.BiPredicate;
import java.util.function.Predicate;
import net.minecraft.datafixer.DataFixTypes;
import net.minecraft.util.math.ChunkPos;
import net.minecraft.world.PersistentState;
import net.minecraft.world.PersistentStateType;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

public class ChunkTicketManager extends PersistentState {
	private static final int DEFAULT_TICKETS_MAP_SIZE = 4;
	private static final Logger LOGGER = LogUtils.getLogger();
	private static final Codec<Pair<ChunkPos, ChunkTicket>> TICKET_POS_CODEC = Codec.mapPair(ChunkPos.CODEC.fieldOf("chunk_pos"), ChunkTicket.CODEC).codec();
	public static final Codec<ChunkTicketManager> CODEC = RecordCodecBuilder.create(
		instance -> instance.group(TICKET_POS_CODEC.listOf().optionalFieldOf("tickets", List.of()).forGetter(ChunkTicketManager::getTickets))
			.apply(instance, ChunkTicketManager::create)
	);
	public static final PersistentStateType<ChunkTicketManager> STATE_TYPE = new PersistentStateType<>(
		"chunks", ChunkTicketManager::new, CODEC, DataFixTypes.SAVED_DATA_FORCED_CHUNKS
	);
	private final Long2ObjectOpenHashMap<List<ChunkTicket>> tickets;
	private final Long2ObjectOpenHashMap<List<ChunkTicket>> savedTickets;
	private LongSet forcedChunks = new LongOpenHashSet();
	@Nullable
	private ChunkTicketManager.LevelUpdater loadingLevelUpdater;
	@Nullable
	private ChunkTicketManager.LevelUpdater simulationLevelUpdater;

	private ChunkTicketManager(Long2ObjectOpenHashMap<List<ChunkTicket>> tickets, Long2ObjectOpenHashMap<List<ChunkTicket>> savedTickets) {
		this.tickets = tickets;
		this.savedTickets = savedTickets;
		this.recomputeForcedChunks();
	}

	public ChunkTicketManager() {
		this(new Long2ObjectOpenHashMap<>(4), new Long2ObjectOpenHashMap<>());
	}

	private static ChunkTicketManager create(List<Pair<ChunkPos, ChunkTicket>> tickets) {
		Long2ObjectOpenHashMap<List<ChunkTicket>> long2ObjectOpenHashMap = new Long2ObjectOpenHashMap<>();

		for (Pair<ChunkPos, ChunkTicket> pair : tickets) {
			ChunkPos chunkPos = pair.getFirst();
			List<ChunkTicket> list = long2ObjectOpenHashMap.computeIfAbsent(
				chunkPos.toLong(), (Long2ObjectFunction<? extends List<ChunkTicket>>)(l -> new ObjectArrayList<>(4))
			);
			list.add(pair.getSecond());
		}

		return new ChunkTicketManager(new Long2ObjectOpenHashMap<>(4), long2ObjectOpenHashMap);
	}

	private List<Pair<ChunkPos, ChunkTicket>> getTickets() {
		List<Pair<ChunkPos, ChunkTicket>> list = new ArrayList();
		this.forEachTicket((pos, ticket) -> {
			if (ticket.getType().persist()) {
				list.add(new Pair<>(pos, ticket));
			}
		});
		return list;
	}

	private void forEachTicket(BiConsumer<ChunkPos, ChunkTicket> ticketConsumer) {
		forEachTicket(ticketConsumer, this.tickets);
		forEachTicket(ticketConsumer, this.savedTickets);
	}

	private static void forEachTicket(BiConsumer<ChunkPos, ChunkTicket> ticketConsumer, Long2ObjectOpenHashMap<List<ChunkTicket>> tickets) {
		for (Entry<List<ChunkTicket>> entry : Long2ObjectMaps.fastIterable(tickets)) {
			ChunkPos chunkPos = new ChunkPos(entry.getLongKey());

			for (ChunkTicket chunkTicket : (List)entry.getValue()) {
				ticketConsumer.accept(chunkPos, chunkTicket);
			}
		}
	}

	public void promoteToRealTickets() {
		for (Entry<List<ChunkTicket>> entry : Long2ObjectMaps.fastIterable(this.savedTickets)) {
			for (ChunkTicket chunkTicket : (List)entry.getValue()) {
				this.addTicket(entry.getLongKey(), chunkTicket);
			}
		}

		this.savedTickets.clear();
	}

	public void setLoadingLevelUpdater(@Nullable ChunkTicketManager.LevelUpdater loadingLevelUpdater) {
		this.loadingLevelUpdater = loadingLevelUpdater;
	}

	public void setSimulationLevelUpdater(@Nullable ChunkTicketManager.LevelUpdater simulationLevelUpdater) {
		this.simulationLevelUpdater = simulationLevelUpdater;
	}

	public boolean hasTickets() {
		return !this.tickets.isEmpty();
	}

	public List<ChunkTicket> getTickets(long pos) {
		return this.tickets.getOrDefault(pos, List.of());
	}

	private List<ChunkTicket> getTicketsMutable(long pos) {
		return this.tickets.computeIfAbsent(pos, (Long2ObjectFunction<? extends List<ChunkTicket>>)(chunkPos -> new ObjectArrayList<>(4)));
	}

	public void addTicket(ChunkTicketType type, ChunkPos pos, int radius) {
		ChunkTicket chunkTicket = new ChunkTicket(type, ChunkLevels.getLevelFromType(ChunkLevelType.FULL) - radius);
		this.addTicket(pos.toLong(), chunkTicket);
	}

	public void addTicket(ChunkTicket ticket, ChunkPos pos) {
		this.addTicket(pos.toLong(), ticket);
	}

	public boolean addTicket(long pos, ChunkTicket ticket) {
		List<ChunkTicket> list = this.getTicketsMutable(pos);

		for (ChunkTicket chunkTicket : list) {
			if (ticketsEqual(ticket, chunkTicket)) {
				chunkTicket.refreshExpiry();
				this.markDirty();
				return false;
			}
		}

		int i = getLevel(list, true);
		int j = getLevel(list, false);
		list.add(ticket);
		if (ticket.getType().isForSimulation() && ticket.getLevel() < i && this.simulationLevelUpdater != null) {
			this.simulationLevelUpdater.update(pos, ticket.getLevel(), true);
		}

		if (ticket.getType().isForLoading() && ticket.getLevel() < j && this.loadingLevelUpdater != null) {
			this.loadingLevelUpdater.update(pos, ticket.getLevel(), true);
		}

		if (ticket.getType().equals(ChunkTicketType.FORCED)) {
			this.forcedChunks.add(pos);
		}

		this.markDirty();
		return true;
	}

	private static boolean ticketsEqual(ChunkTicket a, ChunkTicket b) {
		return b.getType() == a.getType() && b.getLevel() == a.getLevel();
	}

	public int getLevel(long pos, boolean forSimulation) {
		return getLevel(this.getTickets(pos), forSimulation);
	}

	private static int getLevel(List<ChunkTicket> tickets, boolean forSimulation) {
		ChunkTicket chunkTicket = getActiveTicket(tickets, forSimulation);
		return chunkTicket == null ? ChunkLevels.INACCESSIBLE + 1 : chunkTicket.getLevel();
	}

	@Nullable
	private static ChunkTicket getActiveTicket(@Nullable List<ChunkTicket> tickets, boolean forSimulation) {
		if (tickets == null) {
			return null;
		} else {
			ChunkTicket chunkTicket = null;

			for (ChunkTicket chunkTicket2 : tickets) {
				if (chunkTicket == null || chunkTicket2.getLevel() < chunkTicket.getLevel()) {
					if (forSimulation && chunkTicket2.getType().isForSimulation()) {
						chunkTicket = chunkTicket2;
					} else if (!forSimulation && chunkTicket2.getType().isForLoading()) {
						chunkTicket = chunkTicket2;
					}
				}
			}

			return chunkTicket;
		}
	}

	public void removeTicket(ChunkTicketType type, ChunkPos pos, int radius) {
		ChunkTicket chunkTicket = new ChunkTicket(type, ChunkLevels.getLevelFromType(ChunkLevelType.FULL) - radius);
		this.removeTicket(pos.toLong(), chunkTicket);
	}

	public void removeTicket(ChunkTicket ticket, ChunkPos pos) {
		this.removeTicket(pos.toLong(), ticket);
	}

	public boolean removeTicket(long pos, ChunkTicket ticket) {
		List<ChunkTicket> list = this.tickets.get(pos);
		if (list == null) {
			return false;
		} else {
			boolean bl = false;
			Iterator<ChunkTicket> iterator = list.iterator();

			while (iterator.hasNext()) {
				ChunkTicket chunkTicket = (ChunkTicket)iterator.next();
				if (ticketsEqual(ticket, chunkTicket)) {
					iterator.remove();
					bl = true;
					break;
				}
			}

			if (!bl) {
				return false;
			} else {
				if (list.isEmpty()) {
					this.tickets.remove(pos);
				}

				if (ticket.getType().isForSimulation() && this.simulationLevelUpdater != null) {
					this.simulationLevelUpdater.update(pos, getLevel(list, true), false);
				}

				if (ticket.getType().isForLoading() && this.loadingLevelUpdater != null) {
					this.loadingLevelUpdater.update(pos, getLevel(list, false), false);
				}

				if (ticket.getType().equals(ChunkTicketType.FORCED)) {
					this.recomputeForcedChunks();
				}

				this.markDirty();
				return true;
			}
		}
	}

	private void recomputeForcedChunks() {
		this.forcedChunks = this.getAllChunksMatching(ticket -> ticket.getType().equals(ChunkTicketType.FORCED));
	}

	public String getDebugString(long pos, boolean forSimulation) {
		List<ChunkTicket> list = this.getTickets(pos);
		ChunkTicket chunkTicket = getActiveTicket(list, forSimulation);
		return chunkTicket == null ? "no_ticket" : chunkTicket.toString();
	}

	public void tick(ServerChunkLoadingManager serverChunkLoadingManager) {
		this.removeTicketsIf((long_, chunkTicket) -> {
			ChunkHolder chunkHolder = serverChunkLoadingManager.getCurrentChunkHolder(long_);
			boolean bl = chunkHolder != null && !chunkHolder.isSavable() && chunkTicket.getType().isForSimulation();
			if (bl) {
				return false;
			} else {
				chunkTicket.tick();
				return chunkTicket.isExpired();
			}
		}, null);
		this.markDirty();
	}

	public void shutdown() {
		this.removeTicketsIf((long_, chunkTicket) -> chunkTicket.getType() != ChunkTicketType.UNKNOWN, this.savedTickets);
	}

	public void removeTicketsIf(BiPredicate<Long, ChunkTicket> biPredicate, @Nullable Long2ObjectOpenHashMap<List<ChunkTicket>> transferTo) {
		ObjectIterator<Entry<List<ChunkTicket>>> objectIterator = this.tickets.long2ObjectEntrySet().fastIterator();
		boolean bl = false;

		while (objectIterator.hasNext()) {
			Entry<List<ChunkTicket>> entry = (Entry<List<ChunkTicket>>)objectIterator.next();
			Iterator<ChunkTicket> iterator = ((List)entry.getValue()).iterator();
			long l = entry.getLongKey();
			boolean bl2 = false;
			boolean bl3 = false;

			while (iterator.hasNext()) {
				ChunkTicket chunkTicket = (ChunkTicket)iterator.next();
				if (biPredicate.test(l, chunkTicket)) {
					if (transferTo != null) {
						List<ChunkTicket> list = transferTo.computeIfAbsent(
							l, (Long2ObjectFunction<? extends List<ChunkTicket>>)(pos -> new ObjectArrayList<>(((List)entry.getValue()).size()))
						);
						list.add(chunkTicket);
					}

					iterator.remove();
					if (chunkTicket.getType().isForLoading()) {
						bl3 = true;
					}

					if (chunkTicket.getType().isForSimulation()) {
						bl2 = true;
					}

					if (chunkTicket.getType().equals(ChunkTicketType.FORCED)) {
						bl = true;
					}
				}
			}

			if (bl3 || bl2) {
				if (bl3 && this.loadingLevelUpdater != null) {
					this.loadingLevelUpdater.update(l, getLevel((List<ChunkTicket>)entry.getValue(), false), false);
				}

				if (bl2 && this.simulationLevelUpdater != null) {
					this.simulationLevelUpdater.update(l, getLevel((List<ChunkTicket>)entry.getValue(), true), false);
				}

				this.markDirty();
				if (((List)entry.getValue()).isEmpty()) {
					objectIterator.remove();
				}
			}
		}

		if (bl) {
			this.recomputeForcedChunks();
		}
	}

	public void updateLevel(int level, ChunkTicketType type) {
		List<Pair<ChunkTicket, Long>> list = new ArrayList();

		for (Entry<List<ChunkTicket>> entry : this.tickets.long2ObjectEntrySet()) {
			for (ChunkTicket chunkTicket : (List)entry.getValue()) {
				if (chunkTicket.getType() == type) {
					list.add(Pair.of(chunkTicket, entry.getLongKey()));
				}
			}
		}

		for (Pair<ChunkTicket, Long> pair : list) {
			Long long_ = pair.getSecond();
			ChunkTicket chunkTicketx = pair.getFirst();
			this.removeTicket(long_, chunkTicketx);
			ChunkTicketType chunkTicketType = chunkTicketx.getType();
			this.addTicket(long_, new ChunkTicket(chunkTicketType, level));
		}
	}

	public boolean setChunkForced(ChunkPos pos, boolean forced) {
		ChunkTicket chunkTicket = new ChunkTicket(ChunkTicketType.FORCED, ServerChunkLoadingManager.FORCED_CHUNK_LEVEL);
		return forced ? this.addTicket(pos.toLong(), chunkTicket) : this.removeTicket(pos.toLong(), chunkTicket);
	}

	public LongSet getForcedChunks() {
		return this.forcedChunks;
	}

	private LongSet getAllChunksMatching(Predicate<ChunkTicket> predicate) {
		LongOpenHashSet longOpenHashSet = new LongOpenHashSet();

		for (Entry<List<ChunkTicket>> entry : Long2ObjectMaps.fastIterable(this.tickets)) {
			for (ChunkTicket chunkTicket : (List)entry.getValue()) {
				if (predicate.test(chunkTicket)) {
					longOpenHashSet.add(entry.getLongKey());
					break;
				}
			}
		}

		return longOpenHashSet;
	}

	@FunctionalInterface
	public interface LevelUpdater {
		void update(long pos, int level, boolean added);
	}
}
